#!/bin/sh

# Function to log messages (optional, can be adapted from S41network_mac if needed)
log_message() {
    SCRIPT_NAME=$(basename "$0")
    echo "${SCRIPT_NAME}: $1"
    # logger -t "${SCRIPT_NAME}" "$1" # Uncomment if logger is available and preferred
}

# 通用 UBI 卷挂载函数
# 支持两种类型：
# - dynamic: UBI 动态卷，挂载为 ubifs（可读写）
# - static: UBI 静态卷，挂载为 squashfs（只读）
function mount_ubi_volume() {
    local vol_type=$1      # "dynamic" 或 "static"
    local mtd_num=$2
    local mount_point=$3
    local vol_name=$4
    local ubi_num=""

    log_message "Processing mtd${mtd_num} for ${mount_point}..."

    # 尝试附加 MTD 到 UBI（如果已附加会失败，但我们会从现有设备中查找）
    local ubi_output=$(ubiattach /dev/ubi_ctrl -m ${mtd_num} 2>&1)
    local attach_result=$?

    if [ ${attach_result} -eq 0 ]; then
        # 附加成功，从输出提取 UBI 设备编号
        ubi_num=$(echo "${ubi_output}" | sed -n 's/.*UBI device number \([0-9]\+\).*/\1/p' | head -1)
        if [ -n "${ubi_num}" ]; then
            log_message "MTD${mtd_num} attached as /dev/ubi${ubi_num}"
        fi
    else
        # 附加失败（可能已经附加），从 dmesg 或 sysfs 查找
        log_message "ubiattach returned error (may already be attached), searching for existing device..."
    fi

    # 如果从 ubiattach 输出没有获取到编号，尝试其他方法
    if [ -z "${ubi_num}" ]; then
        # 方法1: 从 dmesg 查找
        ubi_num=$(dmesg | grep "ubi.*: attaching mtd${mtd_num}" | tail -1 | sed -n 's/^ubi\([0-9]\+\):.*/\1/p')
    fi

    if [ -z "${ubi_num}" ]; then
        # 方法2: 遍历 /sys/class/ubi/ 查找对应的 MTD
        local sys_ubi_dir="/sys/class/ubi"
        if [ -d "${sys_ubi_dir}" ]; then
            for ubi_device in ${sys_ubi_dir}/ubi[0-9]*; do
                if [ -d "${ubi_device}" ] && [ -f "${ubi_device}/mtd_num" ]; then
                    local mtd_check=$(cat "${ubi_device}/mtd_num")
                    if [ "${mtd_check}" = "${mtd_num}" ]; then
                        ubi_num=$(basename "${ubi_device}" | sed 's/ubi//')
                        log_message "Found existing UBI device: /dev/ubi${ubi_num} for mtd${mtd_num}"
                        break
                    fi
                fi
            done
        fi
    fi

    if [ -z "${ubi_num}" ]; then
        log_message "Failed to determine UBI device number for mtd${mtd_num}"
        return 1
    fi

    # 根据卷类型选择不同的挂载方式
    if [ "${vol_type}" = "static" ]; then
        # 静态卷：SquashFS，需要创建 ubiblock 设备
        log_message "Mounting static volume (SquashFS) from /dev/ubi${ubi_num}_0..."

        # 检查 ubiblock 设备是否已经存在
        if [ ! -e "/dev/ubiblock${ubi_num}_0" ]; then
            log_message "Creating ubiblock device for ubi${ubi_num}_0..."
            ubiblock --create /dev/ubi${ubi_num}_0
            if [ $? -ne 0 ]; then
                log_message "Failed to create ubiblock device for ubi${ubi_num}_0"
                return 1
            fi
        else
            log_message "ubiblock device /dev/ubiblock${ubi_num}_0 already exists"
        fi

        # 挂载 SquashFS
        mount -t squashfs /dev/ubiblock${ubi_num}_0 ${mount_point}
        if [ $? -eq 0 ]; then
            log_message "Mount ${mount_point} success (squashfs)"
            return 0
        else
            log_message "Mount ${mount_point} failed (squashfs)"
            return 1
        fi

    elif [ "${vol_type}" = "dynamic" ]; then
        # 动态卷：UBIFS，需要创建 UBI 卷

        # 检查卷是否已经存在
        if [ -e "/dev/ubi${ubi_num}_0" ]; then
            log_message "Volume ${vol_name} already exists on /dev/ubi${ubi_num}"
        else
            log_message "Attempting to create volume ${vol_name} on /dev/ubi${ubi_num}..."
            ubimkvol /dev/ubi${ubi_num} -m -N ${vol_name}
            if [ $? -ne 0 ]; then
                log_message "Warning: ubimkvol failed, but will try to mount anyway"
            fi
        fi

        log_message "Attempting to mount ubi${ubi_num}_0 to ${mount_point}..."
        # 挂载UBI卷
        mount -t ubifs ubi${ubi_num}_0 ${mount_point}

        if [ $? -ne 0 ]; then
            log_message "${mount_point} mount fail! Attempting to reformat and retry..."
            sync

            flash_erase /dev/mtd${mtd_num} 0 0
            if [ $? -ne 0 ]; then
                log_message "flash_erase /dev/mtd${mtd_num} failed!"
                return 1
            fi

            ubidetach -m ${mtd_num}
            if [ $? -ne 0 ]; then
                log_message "ubidetach -m ${mtd_num} failed!"
            fi

            # Re-attach
            ubi_output=$(ubiattach /dev/ubi_ctrl -m ${mtd_num} 2>&1)
            if [ $? -ne 0 ]; then
                log_message "Second ubiattach failed for mtd${mtd_num}!"
                return 1
            fi

            ubi_num=$(echo "${ubi_output}" | sed -n 's/.*UBI device number \([0-9]\+\).*/\1/p' | head -1)
            if [ -z "${ubi_num}" ]; then
                ubi_num=$(dmesg | grep "ubi.*: attaching mtd${mtd_num}" | tail -1 | sed -n 's/^ubi\([0-9]\+\):.*/\1/p')
            fi

            ubimkvol /dev/ubi${ubi_num} -m -N ${vol_name}
            mount -t ubifs ubi${ubi_num}_0 ${mount_point}

            if [ $? -ne 0 ]; then
                log_message "Mount ${mount_point} failed after retry!"
                return 1
            fi
        fi

        log_message "Mount ${mount_point} success (ubifs)"
        return 0
    else
        log_message "Unknown volume type: ${vol_type}"
        return 1
    fi
}

start() {
    log_message "Starting partition mounting..."

    # Mount assets partition (mtd6) - 20M，UBI静态卷，SquashFS文件系统
    mount_ubi_volume "static" 6 "/mnt/assets" "assets"
    log_message "mount /mnt/assets done, result: $?"

    # Mount config partition (mtd7) - 6M，UBI动态卷，UBIFS文件系统
    mount_ubi_volume "dynamic" 7 "/mnt/config" "config"
    log_message "mount /mnt/config done, result: $?"

    # Mount log partition (mtd8) - 9M，UBI动态卷，UBIFS文件系统
    mount_ubi_volume "dynamic" 8 "/mnt/log" "log"
    log_message "mount /mnt/log done, result: $?"

    # Mount data partition (mtd9) - 32M，UBI动态卷，UBIFS文件系统
    mount_ubi_volume "dynamic" 9 "/mnt/data" "data"
    log_message "mount /mnt/data done, result: $?"

    log_message "Partition mounting finished."
    return 0
}

stop() {
    log_message "Stop command received for partition mounting. No specific action taken."
    # Typically, unmounting is handled by shutdown scripts or is not necessary for "stop".
}

case "$1" in
    start)
        start
        ;;
    stop)
        stop
        ;;
    restart|reload)
        log_message "Restart/Reload requested for partition mounting."
        # For mounting, a restart might involve unmounting then mounting,
        # or just ensuring mounts are present. Given this is early boot,
        # simply re-running start logic is often sufficient.
        # If unmounting is desired, it should be implemented carefully.
        # stop # Optional: if unmounting logic is added to stop()
        start
        ;;
    *)
        echo "Usage: $(basename "$0") {start|stop|restart}"
        exit 1
        ;;
esac

exit $? 